#include "asan_alloctor.h"
#include "asan_interceptors.h"
#include "asan_internal.h"
#include "asan_utils.h"
#include "asan_mapping.h"
#include "asan_report.h"
#include "asan_avltree.h"

namespace __sanitizer {

Allocator* allocatorPtr;

AsanChunkTree::AsanChunkTree(){
  OnInit();
}

void FreeSubTree(AsanChunk* chunk) {
    if (chunk->left != nullptr)
        FreeSubTree(chunk->left);
    if (chunk->right != nullptr)
        FreeSubTree(chunk->right);
    REAL(free(chunk));
}

void AsanChunkTree::ClearAllChunks() {
  VReport(ASAN_LOG_DEBUG, "Free all AsanChunks\n");
  if (root != nullptr)
    FreeSubTree(root);
  root = nullptr;
}

AsanChunkTree::~AsanChunkTree() {
  // To not link with libstdc++, the destructor must be empty
  // if (root != nullptr)
  //   ClearAllChunks();
}

AsanChunk* AsanChunkTree::FindChunkWithBeg(AsanChunk* node, uptr beg) const {
  if (node == nullptr)
    return node;
  if (beg > node->end_addr) {
    return FindChunkWithBeg(node->right, beg);
  }
  else if (beg < node->beg_addr) {
    return FindChunkWithBeg(node->left, beg);
  }
  else {
    if (beg == node->beg_addr) {
      return node;
    }
    else {
      // beg is at the middle of a chunk
      return nullptr;
    }
  }
}

AsanChunk* AsanChunkTree::FindChunkContainsAddr(AsanChunk* node, uptr addr) const {
  if (node == nullptr)
    return node;
  if (addr > node->end_addr) {
    return FindChunkContainsAddr(node->right, addr);
  }
  else if (addr < node->beg_addr) {
    return FindChunkContainsAddr(node->left, addr);
  }
  else {
    return node;
  }
}

AsanChunk* AsanChunkTree::GetListLeftAsanChunk(const AsanChunk* chunk) const {
  if (chunk == nullptr) return nullptr;
  if (chunk->left != nullptr) {
    return MaxValueChunk(chunk->left);
  }
  if (chunk->father == nullptr) return nullptr;
  if (chunk->father->right == chunk) return chunk->father;
  // return nullptr;
  AsanChunk* cur = chunk->father;
  AsanChunk* ancestor = cur->father;
  while (ancestor != nullptr && ancestor->left == cur) {
    cur = ancestor;
    ancestor = cur->father;
  }
  return ancestor;
}

AsanChunk* AsanChunkTree::GetListRightAsanChunk(const AsanChunk* chunk) const {
  if (chunk == nullptr) return nullptr;
  if (chunk->right != nullptr) {
    return MinValueChunk(chunk->right);
  }
  if (chunk->father == nullptr) return nullptr;
  if (chunk->father->left == chunk) return chunk->father;
  // return nullptr;
  AsanChunk* cur = chunk->father;
  AsanChunk* ancestor = cur->father;
  while (ancestor != nullptr && ancestor->right == cur) {
    cur = ancestor;
    ancestor = cur->father;
  }
  return ancestor;
}

void AsanChunkTree::InsertChunk(AsanChunk* chunk) {
  root = __sanitizer::InsertChunk(root, chunk);
  root->father = nullptr;
}

void AsanChunkTree::DeleteChunk(AsanChunk* chunk) {
  root = __sanitizer::DeleteChunk(root, chunk->beg_addr);
  if (root != nullptr)
    root->father = nullptr;
}

AsanChunk* AsanChunkTree::CreateLeftRedzoneChunk(const AsanChunk* chunk) {
  AsanChunk* list_left = GetListLeftAsanChunk(chunk);
  if (list_left == nullptr) {
    // As large redzone as we want
    // uptr min_start_addr = RoundDownTo(chunk->beg_addr, ASAN_PAGE_SIZE);
    // if (chunk->beg_addr - min_start_addr < redzone_weak_min_size) {
    //   min_start_addr -= redzone_weak_min_size;
    // }
    uptr min_start_addr = chunk->beg_addr - redzone_weak_min_size;
    AsanChunk* new_chunk = NewAsanChunk(min_start_addr, chunk->beg_addr - 1, kAsanHeapLeftRedzoneMagic);
    // We do not insert redzone chunks
    // this->InsertChunk(this->root, new_chunk);
    return new_chunk;
  }
  else {
    uptr min_start_addr = list_left->end_addr + 1;
    if (chunk->beg_addr - min_start_addr > redzone_weak_min_size) {
      min_start_addr = chunk->beg_addr - redzone_weak_min_size;
    }
    else if (chunk->beg_addr == min_start_addr) {
      return nullptr;
    }
    AsanChunk* new_chunk = NewAsanChunk(min_start_addr, chunk->beg_addr - 1, kAsanHeapLeftRedzoneMagic);
    // We do not insert redzone chunks
    // this->InsertChunk(this->root, new_chunk);
    return new_chunk;
  }
}

AsanChunk* AsanChunkTree::CreateRightRedzoneChunk(const AsanChunk* chunk) {
  AsanChunk* list_right = GetListRightAsanChunk(chunk);
  if (list_right == nullptr) {
    // as large redzone as we want
    uptr min_end_addr = RoundUpTo(chunk->end_addr, ASAN_PAGE_SIZE) - 1;
    if (min_end_addr - chunk->end_addr + 1 < redzone_weak_min_size) {
      min_end_addr = chunk->end_addr + redzone_weak_min_size - 1;
    }
    else {
      min_end_addr = Min<uptr>(min_end_addr, chunk->end_addr + redzone_weak_max_size - 1);
    }
    AsanChunk* new_chunk = NewAsanChunk(chunk->end_addr + 1, min_end_addr, kAsanInternalHeapMagic);
    // this->InsertChunk(this->root, new_chunk);
    return new_chunk;
  }
  else {
    uptr min_end_addr = list_right->beg_addr - 1;
    if (min_end_addr - chunk->end_addr + 1 > redzone_weak_min_size) {
      min_end_addr = chunk->end_addr + redzone_weak_min_size;
    }
    else if (chunk->end_addr == min_end_addr) {
      return nullptr;
    }
    AsanChunk* new_chunk = NewAsanChunk(chunk->end_addr + 1, min_end_addr, kAsanInternalHeapMagic);
    // this->InsertChunk(this->root, new_chunk);
    return new_chunk;
  }
}

void PrintSubTree(const AsanChunk* chunk, int indent) {
  for (int i = 0; i < indent; ++i) {
    Printf(" ");
  }
  if (chunk != nullptr) {
    Printf("%p %p %x\n", chunk, chunk->Beg(), chunk->mode & 0x0ff);
    PrintSubTree(chunk->left, indent + 1);
    PrintSubTree(chunk->right, indent + 1);
  }
  else {
    Printf("NONE\n");
  }
}

void AsanChunkTree::PrintTree() const {
  PrintSubTree(root, 0);
}

void AsanChunkTree::PrintList() const {

}

int Allocator::OnInit(AsanShadowMem *ptr) {
  shadow = ptr;
  tree.OnInit();
  tree.redzone_weak_min_size = asan_flags()->redzone;
  tree.redzone_weak_max_size = asan_flags()->max_redzone;
}

int Allocator::OnDelete(){
  VReport(ASAN_LOG_DEBUG, "Free shadow\n");
  if (shadow != nullptr) {
    REAL(free(shadow));
  }
  tree.ClearAllChunks();
}

int Allocator::DoMalloc(void* addr, uptr size) {
  VReport(ASAN_LOG_DEBUG, "%s:%d %s(addr=%p, size=%lld)\n", __FILE__, __LINE__, __func__, addr, size);
  shadow->UnpoisonMem(addr, size);
  AsanChunk * new_chunk = NewAsanChunk((uptr)addr, (uptr)addr + size - 1, 0);
//   tree.PrepareToInsertChunk(new_chunk);
  tree.InsertChunk(new_chunk);
//   REAL(free(new_chunk));
  AsanChunk* left_redzone = tree.CreateLeftRedzoneChunk(new_chunk);
  if (left_redzone != nullptr) {
    shadow->PoisonMem(left_redzone->Beg(), left_redzone->ChunkSize(), left_redzone->mode);
    REAL(free(left_redzone));
  }
  AsanChunk* right_redzone = tree.CreateRightRedzoneChunk(new_chunk);
  if (right_redzone != nullptr) {
    shadow->PoisonMem(right_redzone->Beg(), right_redzone->ChunkSize(), right_redzone->mode);
    REAL(free(right_redzone));
  }
  return 0;
}

int Allocator::DoAlloca(void* addr, uptr size) {
  // We do not insert a chunk with alloca
  // but insert some redzone at right
  shadow->UnpoisonMem(addr, size);
  void* redzone_beg = (void*)((uptr)addr + size);
  shadow->PoisonMem(redzone_beg, redzone_weak_min_size, kAsanStackRightRedzoneMagic);
}

int Allocator::DoFree(void* addr) {
  if (addr == nullptr) return 0;
  VReport(ASAN_LOG_DEBUG, "%s:%d %s(addr=%p)\n", __FILE__, __LINE__, __func__, addr);
  AsanChunk* free_chunk = tree.FindChunkWithBeg(tree.root, (uptr)addr);
  if (free_chunk == nullptr) {
    GET_CALLER_PC_BP_SP;
    if (shadow->IsFreeAfterUse(addr))
      ReportDoubleFree(pc, bp, sp, (uptr)addr);
    else
      ReportFreeNotMalloced(pc, bp, sp, (uptr)addr);
  }
  else if (free_chunk->IsPoisonedChunk()) {
    // We delete the chunk while free
    ASSERT(false, "We should have delete the chunk.");
    GET_CALLER_PC_BP_SP;
    ReportDoubleFree(pc, bp, sp, (uptr)addr);
  }
  DoQuickFree(free_chunk);
  return 0;
}

int Allocator::DoQuickFree(AsanChunk* chunk) {
  // The delete may change the content of the chunk
  // shadow must read from the chunk before delete it
  shadow->PoisonMem((void*)(chunk->beg_addr), chunk->ChunkSize(), kAsanHeapFreeMagic);
  DeleteChunk(chunk);
//   chunk->mode = kAsanHeapFreeMagic;
}

} // namespace __sanitizer
